Rustで実装するmalloc |
您所在的位置：网站首页 › rust malloc › Rustで実装するmalloc |
Rustで実装するmalloc
|
この記事は、NTT Communications Advent Calendar 2021 21日目の記事です。 はじめにこんにちは、イノベーションセンターの鈴ヶ嶺(@suzu_3_14159265)です。普段は、クラウド・ハイブリッドクラウド・エッジデバイスなどを利用したAI/MLシステムに関する業務に従事しています。本日は、Rustで動的メモリ確保(dynamic memory allocation)のmallocを実装してPythonやvimを動かしてみようという内容をお届けします。 また、去年もRustネタのアドベントカレンダーを書いているのでぜひ見ていただけると嬉しいです！ NTTコミュニケーションズ Advent Calendar 2020 Rustで実装するNetflow Collector 実装するmallocのアルゴリズム今回実装するmallocのアルゴリズムは小さなメモリサイズにはSegregated Free Listを用いて、大きなメモリサイズにはmmapを用いる方針で実装します。まず、最初に動的メモリ確保における重要な課題であるメモリ断片化の概要について説明し、それぞれのアルゴリズムの手法について説明します。 課題：メモリ断片化
上記の図ように、時間が経過するにつれてメモリの解放や格納などが行われた結果として未使用領域が飛び飛びに配置され、合計としての空き容量は存在するが連続して要求されたメモリ領域を提供できない課題をメモリ断片化と言います。このような虫食い状態のメモリ配置では、有効な資源(メモリ)を活用できません。このメモリ断片化の課題を次のSegregated Free Listによって解決します。 Segregated Free List
Segregated Free Listは上記の図のように、複数のリストをそれぞれのメモリサイズごとに管理する方式です。これによりメモリ断片化が起きないbest fitなメモリ割り当てを実現します。また、未使用メモリを探索する計算量もそれぞれの管理するリストからunlinkするだけなのでO(1)で済みます。割り当てるメモリ領域がない場合は、sbrkシステムコールによるヒープ拡張を用いてメモリを確保してListに追加するように実装する必要があります。しかし、このSegregated Free Listのみでは全ての要求されうるメモリサイズのリストを用意する必要があり現実的ではないため、今回は512Byte以上のメモリについては後述するmmapを用いて別管理する方法を取ります。 mmap
今回、512byte以上の大きなメモリサイズについてはmmapで別管理してメモリを確保、解放する方法を採用しました。本来のmmapはファイルをメモリにマップするためのシステムコールですが、ファイルディスクリプタ fd に-1を指定し、無名マッピング MAP_ANONYMOUS を利用することでメモリ確保APIとしても使用可能です。この命令を用いて直接kernelからメモリを取得して割り当てます。 非スレッドセーフまた、今回マルチスレッドには未対応な実装となっているためいくつかのプログラムは普通にSegmentation faultで落ちるので注意してください。glibcなどで実装されている詳しいmallocのアルゴリズムについては知りたい方はDoug Leaのmallocという通称 dlmalloc をおすすめします。1つのファイルに実装がまとめられており比較的読みやすいものになっています。 ちなみに以下は、dlmallocのコンパイル方法です。 wget http://gee.cs.oswego.edu/pub/misc/malloc-2.8.4.c gcc -fPIC -shared -o dlmalloc.so malloc-2.8.4.c Rustで共有ライブラリを作成する方法 wrap jemallocまず、 cargo new --lib malloc-rs で新しくライブラリプロジェクトを作成します。そして以下のように Cargo.toml の crate-type に cdylib を追加しましょう。これで cargo build と実行することで*.so が生成され、Rustを用いて他言語から使用可能なライブラリが作成可能になります。 Cargo.toml [package] name = "malloc-rs" version = "0.1.0" edition = "2021" [lib] crate-type = ["cdylib"] [dependencies] libc = "0.2.112" jemalloc-sys = "0.3.2"試しに中身をjemallocにwrapして置き換えるようにコードを書くと次のようになります。実装する関数はmalloc, free, calloc, reallocとなります。jemallocはメモリ断片化の回避と並行処理に対してスケーラブルに動作するmalloc実装です。ちなみに余談ですがRust1.31.1までmallocの実装にjemallocが標準として使用されていました。しかし、システム言語であるのにシステムのメモリアロケータを使用しないのは不自然な点やバイナリサイズの増加が問題として浮かび上がってきたため現在ではシステムのアロケータが使用されるようになっています。 以下のコードを見ていただいたら分かるように、extern "C" と記述することでC APIとして公開できます。そのままだとmangleされるので #[no_mangele] と付け加えましょう。 実際の動作には、LinuxでLD_PRELOAD という環境変数に共有ライブラリを指定して、動的ライブラリを差し替えることができる仕組みを使います。今回、mallocを呼ぶと標準出力で call malloc と出力されるような実装にしたので実際に動作させて確認してみます。 またここでは、writeシステムコールを使用して出力しています。Rustで一般的に使われる println! やlibcの printf などは内部でmallocを使用しているためSegmentation faultで落ちるので注意してください。 src/lib.rs use libc::size_t; extern crate jemalloc_sys; extern crate libc; #[no_mangle] pub unsafe extern "C" fn malloc(size: size_t) -> *mut libc::c_void { let message = "call malloc\n"; let buf = message.as_ptr() as *const libc::c_void; let buf_len = message.len(); libc::write(1, buf, buf_len); jemalloc_sys::malloc(size) } #[no_mangle] pub unsafe extern "C" fn realloc(p: *mut libc::c_void, size: size_t) -> *mut libc::c_void { jemalloc_sys::realloc(p, size) } #[no_mangle] pub unsafe extern "C" fn calloc(number: size_t, size: size_t) -> *mut libc::c_void { jemalloc_sys::calloc(number, size) } #[no_mangle] pub unsafe extern "C" fn free(p: *mut libc::c_void) { jemalloc_sys::free(p) } 実行結果 cargo new --lib malloc-rs # vim malloc-rs/src/lib.rs cargo build export LD_PRELOAD=`pwd`/target/debug/libmalloc_rs.so ls # 適当なコマンドを実行してみる
Headerはやや冗長ですが、3つのメンバのメモリサイズのsize , mmapされたメモリかis_mmap, 次の要素を指し示す next を設定しました。今回のケースにおいて、mmapされたメモリかどうかはsizeで判断可能(512Byteより大きい)です。また本来であれば、割り当ては8Byteにアライメントされるためsizeの下位3bitに空き容量があります。それも今回は、実装の見通しやすさを優先して次の図のようにしました。
全体的な実装は以下のようになります。 get_headerは、割り当てたメモリのポインタからHeaderサイズを引いてHeaderを取得します。 init_mallocで、Segregated Free Listを初期化します。 add_listは、Segregated Free Listに割り当て可能なメモリがない場合にヒープからsbrkを用いてメモリを確保してリストを追加します find_chunkは、Segregated Free Listからsize(Byte)のメモリを探索します extern crate libc; use libc::{c_void, size_t}; /// malloc Header struct Header { /// size of buffer size: size_t, /// flag of mmap is_mmap: size_t, // next header next: *mut Header, } /// 8byte aligment const ALIGN: usize = 8; /// max byte in segregated free list const MAX_BYTE: usize = 512; /// init size of one free list const INIT_LIST_SIZE: usize = 512; /// add size of free list const ADD_LIST_SIZE: usize = 512; /// number of free list const NUM_LIST: usize = MAX_BYTE / ALIGN + 1; /// init size of heap(sbrk) const INIT_HEAP_SIZE: usize = NUM_LIST * (INIT_LIST_SIZE + std::mem::size_of::()); /// flag of call init_malloc static mut IS_INIT_MALLOC: bool = false; /// segregated free list static mut FREE_LISTS: [*mut Header; (NUM_LIST)] = [std::ptr::null_mut(); (NUM_LIST)]; fn get_align(size: usize) -> usize { (size + ALIGN - 1) / ALIGN * ALIGN } unsafe fn get_header(p: *mut c_void) -> *mut Header { let header = p.sub(std::mem::size_of::()) as *mut Header; header } /// init malloc function /// Setup the initial value of the segregated free list using sbrk from heap. unsafe fn init_malloc() -> Result { IS_INIT_MALLOC = true; let current_p = libc::sbrk(0); let ret = libc::sbrk((INIT_HEAP_SIZE as isize).try_into().unwrap()); if ret != current_p { // fail sbrk return Err(ret); } // init segregated free list let mut p = ret; for i in 1..NUM_LIST { FREE_LISTS[i] = p as *mut Header; let num_header = INIT_LIST_SIZE / (i * ALIGN); for j in 0..num_header { let mut header = p as *mut Header; let size = i * ALIGN; (*header).size = size; (*header).is_mmap = 0; (*header).next = std::ptr::null_mut(); let next_p = p.add(size + std::mem::size_of::()); if j != (num_header - 1) { (*header).next = next_p as *mut Header; } else { // last element (*header).next = std::ptr::null_mut(); } p = next_p; } } Ok(()) } /// add segregated free list /// When there is no more memory in the segregated free list, use sbrk to add memory from the heap. unsafe fn add_list(size: usize) -> Result { let current_p = libc::sbrk(0); let num_header = ADD_LIST_SIZE / size; let ret = libc::sbrk( (num_header * (size + std::mem::size_of::())) .try_into() .unwrap(), ); if ret != current_p { // fail sbrk return Err(ret); } let mut p = ret; for j in 0..num_header { let mut header = p as *mut Header; (*header).size = size; (*header).is_mmap = 0; (*header).next = std::ptr::null_mut(); let next_p = p.add(size + std::mem::size_of::()); if j != (num_header - 1) { (*header).next = next_p as *mut Header; } else { // last element (*header).next = std::ptr::null_mut(); } p = next_p; } Ok(ret as *mut Header) } /// find header function /// Get a header of a given size from the segregated free list. unsafe fn find_chunk(size: usize) -> Result { // index of segregated free list let index = size / 8; if FREE_LISTS[index] == std::ptr::null_mut() { let new_list_ret = add_list(size); match new_list_ret { Ok(new_list) => { FREE_LISTS[index] = new_list; } Err(err) => { return Err(err); } } } let header = FREE_LISTS[index]; // unlink chunk FREE_LISTS[index] = (*header).next; Ok(header) } /// malloc function #[no_mangle] pub unsafe extern "C" fn malloc(size: size_t) -> *mut c_void { if size == 0 { return std::ptr::null_mut(); } if !IS_INIT_MALLOC { if init_malloc().is_err() { return std::ptr::null_mut(); } } let size_align = get_align(size); if size_align *mut c_void { let size_align = get_align(size); if p == std::ptr::null_mut() { return malloc(size_align); } let new_p = malloc(size_align); let header = get_header(p); let memcpy_size = if (*header).size *mut c_void { let new_p = malloc(size * number); libc::memset(new_p, 0, size * number); return new_p; } /// free function #[no_mangle] pub unsafe extern "C" fn free(p: *mut c_void) { if p == std::ptr::null_mut() { return; } let header = get_header(p); let size = (*header).size; if (*header).is_mmap == 1 { // free mmap let munmap_ret = libc::munmap(p.sub(std::mem::size_of::()), size); debug_assert!(munmap_ret == 0); } else { // reuse in segregated free list let index = size / ALIGN; let first_header = FREE_LISTS[index]; FREE_LISTS[index] = header; (*header).next = first_header; } } 実行結果実行には、シングルスレッドで動作するプログラムとしてPythonとvimを実行してみました。 Python(numpy)ここでは以下のように、numpyでランダムに大きなサイズの配列を作成し続ける処理で問題ないかを確認しました。 cargo build export LD_PRELOAD=`pwd`/target/debug/libmalloc_rs.so python3 |
【本文地址】